package com.agent.rag.consumer;

import com.agent.commons.Result.NoticeResult;
import com.agent.commons.enums.FileType;
import com.agent.commons.ws.NoticeWebsocket;
import com.agent.rag.client.MilvusClient;
import com.agent.rag.client.SplitterClient;
import com.agent.rag.entity.FileActionMessage;
import com.agent.rag.entity.MilvusEntity;
import com.agent.rag.utils.FileUtil;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.amqp.core.ExchangeTypes;
import org.springframework.amqp.rabbit.annotation.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.InputStream;
import java.net.MalformedURLException;
import java.net.URISyntaxException;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Rag模块消息队列监听
 */
@Component
@Slf4j
public class RagComsumer {

    /**
     * 文本切割工具
     */
    @Resource
    private SplitterClient split;

    @Autowired
    private EmbeddingModel embeddingModel;
    @Autowired
    private MilvusClient milvusClient;

    /**
     * 文件处理
     * @param file
     */
    @RabbitListener(bindings = @QueueBinding(
            value = @Queue(value ="file.queue",durable ="true"),
            exchange=@Exchange(value="rag.file.direct",type= ExchangeTypes.DIRECT),
            key = {"action"},
            arguments = @Argument(name = "x-delay",value = "lazy")
    ))
    public void fileAction(FileActionMessage file) throws MalformedURLException, URISyntaxException {
        List<MilvusEntity> milvusEntities=new ArrayList<>();
        System.out.println("文件处理："+file.getUrl());
        URL fileUrl =new URL(file.getUrl());
        try(InputStream inputStream = fileUrl.openStream()){
        switch (Objects.requireNonNull(FileType.getType(file.getType()))) {
            case DOC -> {
            }
            case DOCX -> {
                //doc文件处理
                String convert = FileUtil.docxConvert(inputStream);
                List<String> strings = split.spiltText(convert);
                for(String s : strings){
                    List< String> picture =new ArrayList<>();
                    //构建MilvusEntity
                    MilvusEntity milvusEntity = new MilvusEntity();
                    milvusEntity.setUserId(file.getUserId());
                    //获取每段文本块中的图片UUID
                    picture=getUUID(s);
                    s=dataWash(s);
                    //文本向量化
                    float[] embed = embeddingModel.embed(s);
                    milvusEntity.setText(s);
                    milvusEntity.setPhoto(picture);
                    milvusEntity.setDense_vector(embed);
                    milvusEntities.add(milvusEntity);
                }
            }
            case PDF -> {
                //TODO 添加pdf处理逻辑
                System.out.println("pdf");
            }
            case TXT -> {
                //TODO 添加txt处理逻辑
                System.out.println("txt");
            }
        }
        //数据插入
        milvusClient.insert(milvusEntities);
        NoticeWebsocket.sendMessageByUserId(file.getUserId(), NoticeResult.success("文件导入成功!"));
        }catch (Exception e){
            e.printStackTrace();
        }
    }

    /**
     * 使用正则表达式获取图片UUID
     * @return
     */
    private List<String> getUUID(String text){
        List<String> uuids = new ArrayList<>();
        Pattern pattern = Pattern.compile("[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}");
        Matcher matcher = pattern.matcher(text);
        while (matcher.find()) {
            String uuid = matcher.group();
            uuids.add(uuid);
        }
        return uuids;
    }

    /**
     * 数据清洗
     * @param text
     * @return
     */
    private String dataWash(String text){
        String regex=Pattern.quote("![图片](")+"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}http:[^)]*\\)";
        return text.replaceAll(regex, "");
    }

}
